/*
  Copyright (c) 2010
  The Regents of the University of Michigan
  All Rights Reserved

  Permission is granted to use, copy, create derivative works, and
  redistribute this software and such derivative works for any purpose,
  so long as the name of the University of Michigan is not used in
  any advertising or publicity pertaining to the use or distribution
  of this software without specific, written prior authorization. If
  the above copyright notice or any other identification of the
  University of Michigan is included in any copy of any portion of
  this software, then the disclaimer below must also be included.

  This software is provided as is, without representation or warranty
  of any kind either express or implied, including without limitation
  the implied warranties of merchantability, fitness for a particular
  purpose, or noninfringement.  The Regents of the University of
  Michigan shall not be liable for any damages, including special,
  indirect, incidental, or consequential damages, with respect to any
  claim arising out of or in connection with the use of the software,
  even if it has been or is hereafter advised of the possibility of
  such damages.
*/

#include <arpa/inet.h>
#include <assert.h>
#include <byteswap.h>
#include <endian.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <glib.h>

#include "handlers.h"
#include "id.h"
#include "serialization.h"
#include "utils.h"


/* Extract a value of the given type into the lvalue target. Requires that
   packet is a char buf holding the value to extract in network byte order. */
#define DECODE_PRIMITIVE_TYPE(target)  do {      \
    typedef typeof(target) type;                 \
    type *ptr = (type *) packet;                 \
    packet += sizeof(type);                      \
    if (G_UNLIKELY(packet > end)) {              \
      goto ERROR;                                \
    }                                            \
    (target) = *ptr;                             \
    if (sizeof(target) == 8) {                   \
      (target) = bswap_64(target);               \
    }                                            \
    else if (sizeof(target) == 4) {              \
      (target) = ntohl(target);                  \
    }                                            \
    else if (sizeof(target) == 2) {              \
      (target) = ntohs(target);                  \
    }                                            \
    else if (sizeof(target) != 1) {              \
      assert(0);                                 \
    }                                            \
  } while(0)


static const char* decode_contact(const char *packet,
                                  const char *end,
                                  struct ContactAddress *out);


/* Decode a ContactAddress from packet into out_addr. Return the
   number of bytes decoded, or negative on error. Caller is
   responsible for length checking.*/
static int decode_address(const char *packet, const char *end,
                          struct ContactAddress *out_addr) {
  /* Expected format: 1-byte address length (error if not 4.)
     length bytes of address, network byte order.
     2 bytes of port, network byte order. */

  const char *start = packet;
  uint8_t tmp;

  DECODE_PRIMITIVE_TYPE(tmp);
  if (tmp != 4) {
    goto ERROR;
  }

  DECODE_PRIMITIVE_TYPE(out_addr->addr.s_addr);
  /* We'll keep the address in network byte order. */
  out_addr->addr.s_addr = htonl(out_addr->addr.s_addr);

  DECODE_PRIMITIVE_TYPE(out_addr->port);

  return packet - start;

ERROR:
  return -1;
};


static int decode_request_header(const char *packet, const char *end,
                                 struct RequestPacketHeader *out) {
  const char *start = packet;
  int offset;

  assert(out);
  if (end - start < REQUEST_HEADER_SIZE) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Ignoring request with too-short header\n");
#endif
    return -1;
  }

  DECODE_PRIMITIVE_TYPE(out->connection_id);

  DECODE_PRIMITIVE_TYPE(out->action);

  DECODE_PRIMITIVE_TYPE(out->transaction_id);
  
  DECODE_PRIMITIVE_TYPE(out->protocol_version);

  if (out->protocol_version < MIN_SUPPORTED_PROTOCOL_VERSION ||
      out->protocol_version > SUPPORTED_PROTOCOL_VERSION) {
    #ifdef VERBOSE_LOG
    fprintf(stderr, "Ignoring request with unsupported protocol version %d"
            " (shouldn't happen!)\n", out->protocol_version);
    #endif
    return -1;
  }

  // NOTE: we might drop this decode if perf demands later.
  DECODE_PRIMITIVE_TYPE(out->vendor_id);

  DECODE_PRIMITIVE_TYPE(out->network_id);

  // NOTE: we might drop this decode if perf demands later.
  DECODE_PRIMITIVE_TYPE(out->local_protocol_version);

  offset = decode_address(packet, end, &out->node_address);
  if (offset < 0) {
    return offset;
  }

  packet += offset;

  DECODE_PRIMITIVE_TYPE(out->instance_id);

  DECODE_PRIMITIVE_TYPE(out->time);

  return packet-start;

ERROR:
  return -1;
}


static int decode_key(const char *packet, const char *end,
                             char **out) {
  uint8_t length;

  if (end - packet < KEY_SIZE + 1) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Key field too short!\n");
#endif
    return -1;
  }

  DECODE_PRIMITIVE_TYPE(length);
  if (length != KEY_SIZE) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Key with length %d found, expected %d!\n", length,
            KEY_SIZE);
#endif
    return -1;
  }

  *out = clearview_id_new();
  memcpy(*out, packet, KEY_SIZE);
  return KEY_SIZE+1;

ERROR:
  return -1;
}


static int decode_value(const char *packet, const char *end,
                        struct TransportValue *out) {
  const char *start = packet;
  DECODE_PRIMITIVE_TYPE(out->version);

  // Creation time ignored because we don't handle clock skew.
  packet += sizeof(uint64_t);

  DECODE_PRIMITIVE_TYPE(out->value_bytes_count);
  
  if (out->value_bytes_count) {
    if (G_UNLIKELY(packet + out->value_bytes_count > end)) {
#ifdef VERBOSE_LOG
      fprintf(stderr, "Overflow decoding value!\n");
#endif
      goto ERROR;
    }
    out->value_bytes = g_memdup(packet, out->value_bytes_count);
    packet += out->value_bytes_count;
  }

  packet = decode_contact(packet, end, &out->originator);
  if (!packet) {
    goto ERROR;
  }

  // Flags ignored.
  packet += sizeof(uint8_t);

  return packet - start;

ERROR:
  transport_value_free_data(out);
  return -1;
}


static inline int decode_ping_request(const struct RequestPacketHeader *header,
                                      const char *packet, const char *end,
                                      struct ClearviewInstance *cv,
                                      struct ContactAddress *contact) {
  struct RequestPingPacket pkt = {*header};
  handle_ping_request(&pkt, cv, contact);
  return 0;
}


void store_request_free_data(struct RequestStorePacket *req) {
  int ii;

  g_return_if_fail(req);

  for (ii = 0; ii != req->keys_count; ++ii) {
    clearview_id_free(req->keys[ii]);
  }
  g_free(req->keys);

  for (ii = 0; ii != req->value_groups_count; ++ii) {
    transport_value_group_cleanup(&req->value_groups[ii]);
  }

  g_free(req->value_groups);
}


static int decode_store_request(const struct RequestPacketHeader *header,
                                const char *packet, const char *end,
                                struct ClearviewInstance *cv,
                                struct ContactAddress *contact) {
  struct RequestStorePacket pkt = {*header};
  int ii;

  if (end - packet < 5) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Bad store request received!\n");
#endif
    return -1;
  }
  
  DECODE_PRIMITIVE_TYPE(pkt.spoof_id);
    
  DECODE_PRIMITIVE_TYPE(pkt.keys_count);
  if (pkt.keys_count > MAX_KEYS_PER_PACKET) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Store request with too many (%d) keys received!\n",
            pkt.keys_count);
#endif
    return -1;
  }

  pkt.keys = g_malloc0(sizeof(*pkt.keys) * pkt.keys_count);

  for (ii = 0; ii != pkt.keys_count; ++ii) {
    int result = decode_key(packet, end, &pkt.keys[ii]);
    if (result < 0) {
#ifdef VERBOSE_LOG
      fprintf(stderr, "Bad key encountered!\n");
#endif
      store_request_free_data(&pkt);
      return result;
    } else {
      packet += result;
    }
  }

  DECODE_PRIMITIVE_TYPE(pkt.value_groups_count);

  if (pkt.value_groups_count > MAX_KEYS_PER_PACKET) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Store request with too many (%d) val grps received!\n",
            pkt.value_groups_count);
#endif
    goto ERROR;
  }

  pkt.value_groups = g_malloc0(sizeof(*pkt.value_groups) * pkt.value_groups_count);
  for (ii = 0; ii != pkt.value_groups_count; ++ii) {
    uint16_t values_count;
    int jj;
    DECODE_PRIMITIVE_TYPE(values_count);
    transport_value_group_init(&pkt.value_groups[ii], values_count);
    for (jj = 0; jj != values_count; ++jj) {
      int result = decode_value(packet, end, &pkt.value_groups[ii].values[jj]);
      if (result < 0) {
#ifdef VERBOSE_LOG
        fprintf(stderr, "Failed to decode value!\n");
        fprintf(stderr, "%dth of %d values in %dth of %d value groups.\n",
                (jj+1), values_count, (ii+1), pkt.value_groups_count);
#endif
        store_request_free_data(&pkt);
        return result;
      } else {
        packet += result;
      }
    }
  }


  handle_store_request(&pkt, cv, contact);
  return 0;

ERROR:
  store_request_free_data(&pkt);
  return -1;
}

static inline int decode_find_node_request(
    const struct RequestPacketHeader *header,
    const char *packet, const char *end,
    struct ClearviewInstance *cv, struct ContactAddress *contact) {
  struct RequestFindNodePacket pkt = {*header};
  // We don't care about the fields of the request right now.
  // TODO(swolchok): actually decode the packet?
  handle_find_node_request(&pkt, cv, contact);
  return 0;
}


static inline int decode_find_value_request(
    const struct RequestPacketHeader *header,
    const char *packet, const char *end,
    struct ClearviewInstance *cv, struct ContactAddress *contact) {
  struct RequestFindValuePacket pkt = {*header};
  // We don't care about the fields of the request right now.
  // TODO(swolchok): actually decode the packet?
  handle_find_value_request(&pkt, cv, contact);
  return 0;
}


int decode_request(const char *packet, size_t len,
                   struct ClearviewInstance *cv,
                   struct ContactAddress *contact) {
  struct RequestPacketHeader header;
  const char *end = packet + len;

  RETURN_IF_NEGATIVE(decode_request_header(packet, end, &header));
  packet += REQUEST_HEADER_SIZE;

  switch(header.action) {
  case ACT_REQUEST_PING:
    return decode_ping_request(&header, packet, end, cv, contact);
  case ACT_REQUEST_STORE:
    return decode_store_request(&header, packet, end, cv, contact);
  case ACT_REQUEST_FIND_NODE:
    return decode_find_node_request(&header, packet, end, cv, contact);
  case ACT_REQUEST_FIND_VALUE:
    return decode_find_value_request(&header, packet, end, cv, contact);

  default:
    fprintf(stderr, "Ignoring request with action %d\n", header.action);
    break;
  }
  return 0;
}


static int decode_reply_header(const char *packet, const char *end,
                               struct ReplyPacketHeader *out) {
  assert(out);

  if(end - packet < REPLY_HEADER_SIZE) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Ignoring reply with too-short header\n");
#endif
    return -1;
  }

  DECODE_PRIMITIVE_TYPE(out->action);

  DECODE_PRIMITIVE_TYPE(out->transaction_id);

  DECODE_PRIMITIVE_TYPE(out->connection_id);

  DECODE_PRIMITIVE_TYPE(out->protocol_version);
  if (out->protocol_version < MIN_SUPPORTED_PROTOCOL_VERSION ||
      out->protocol_version > SUPPORTED_PROTOCOL_VERSION) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Got packet with protocol version %d, expected %d\n",
            out->protocol_version, SUPPORTED_PROTOCOL_VERSION);
#endif
    return -1;
  }

  DECODE_PRIMITIVE_TYPE(out->vendor_id);

  DECODE_PRIMITIVE_TYPE(out->network_id);

  DECODE_PRIMITIVE_TYPE(out->instance_id);

  return 0;

ERROR:
  return -1;
}


static int decode_ping_reply(const struct ReplyPacketHeader *header,
                             const char *packet, const char *end,
                             struct ClearviewInstance *cv,
                             struct ContactAddress *contact) {
  struct ReplyPingPacket reply = {*header};
  // We ignore Vivaldi, so don't bother to decode it.

  handle_ping_reply(&reply, cv, contact);
  return 0;
}


/* Pre-condition: packet points to a Vivaldi coordinates section of a packet.
   Post-condition: returned pointer points past the Vivaldi coordinates, or it == end. */
static inline const char *skip_vivaldi(const char *packet, const char *end) {
  uint8_t length;
  int ii;

  DECODE_PRIMITIVE_TYPE(length);
  for (ii = 0; ii < length; ++ii) {
    uint8_t coord_type, size;
    DECODE_PRIMITIVE_TYPE(coord_type);
    DECODE_PRIMITIVE_TYPE(size);
    packet += size;
  }

  return packet;
ERROR:
  return end;
}


static const char* decode_contact_address(const char *packet,
    const char *end, struct ContactAddress *out) {
  uint8_t length;
  const char *start = packet;

  DECODE_PRIMITIVE_TYPE(length);
  if (length != 4) {
#ifdef VERBOSE_LOG
    if (length == 6) {
      fprintf(stderr, "Someone's using IPv6. What the heck?\n");
    } else {
      fprintf(stderr, "Error decoding contact address -- bad length %d\n",
              (int)length);
    }
#endif
    goto ERROR;
  }

  assert((packet - start) == 1);
  DECODE_PRIMITIVE_TYPE(out->addr.s_addr);
  /* We want to keep the IP address in network byte order. */
  out->addr.s_addr = htonl(out->addr.s_addr);

  DECODE_PRIMITIVE_TYPE(out->port);

  assert(packet <= end);
  assert((packet - start) == CONTACT_ADDRESS_NET_SIZE);
  return packet;

ERROR:
  return NULL;
}


// We don't care about the extra data from Contacts.
static const char *decode_contact(const char *packet, const char *end,
                                  struct ContactAddress *out) {
  const char *start = packet;
  uint8_t contact_type, contact_protocol_version;

  DECODE_PRIMITIVE_TYPE(contact_type);
  if (contact_type != 1) {
#ifdef VERBOSE_LOG
    fprintf(stderr, "Received contact with bogus type %d\n",
            (int)contact_type);
    fprintf(stderr, "Next bytes: %x %x %x %x\n", packet[0], packet[1], packet[2], packet[3]);
#endif
    goto ERROR;
  }

  DECODE_PRIMITIVE_TYPE(contact_protocol_version);
  packet = decode_contact_address(packet, end, out);
  if (!packet) {
    goto ERROR;
  }

  assert(packet <= end);
  assert(packet - start == CONTACT_NET_SIZE);

  return packet;
ERROR:
  return NULL;
}

static int decode_find_node_reply(const struct ReplyPacketHeader *header,
                                  const char *packet, const char *end,
                                  struct ClearviewInstance *cv,
                                  struct ContactAddress *contact) {
  struct ReplyFindNodePacket reply = {*header};
  uint32_t temp;
  int ii;

  reply.contacts = NULL;
  reply.contacts_count = 0;

  DECODE_PRIMITIVE_TYPE(reply.spoof_id);

  DECODE_PRIMITIVE_TYPE(temp);
  reply.node_type = temp;

  DECODE_PRIMITIVE_TYPE(reply.dht_size);

  packet = skip_vivaldi(packet, end);
  if (packet >= end) { 
    goto ERROR;
  }

  DECODE_PRIMITIVE_TYPE(reply.contacts_count);
  reply.contacts = g_malloc0(reply.contacts_count * sizeof(*reply.contacts));

  for(ii = 0; ii != reply.contacts_count; ++ii) {
    const char *new_val;
    reply.contacts[ii] = contact_new();
    contact_init(reply.contacts[ii]);
    new_val = decode_contact(packet, end, reply.contacts[ii]);
    if (!new_val) {
      goto ERROR;
    }
    g_assert(new_val - packet == CONTACT_NET_SIZE);
    packet = new_val;
  }

  handle_find_node_reply(&reply, cv, contact);
  return 0;

ERROR:
  if (reply.contacts) {
    for (ii = 0; ii != reply.contacts_count; ++ii) {
      contact_free(reply.contacts[ii]);
    }
    g_free(reply.contacts);
  }
  return -1;
}


static int decode_error_reply(const struct ReplyPacketHeader *header,
                              const char *packet, const char *end,
                              struct ClearviewInstance *cv,
                              struct ContactAddress *contact) {
  struct ReplyErrorPacket pkt = {*header};

  DECODE_PRIMITIVE_TYPE(pkt.error_type);

  if (pkt.error_type == WRONG_ADDRESS) {
    packet = decode_contact_address(packet, end, &pkt.sender_address);
  }

  handle_error_reply(&pkt, cv, contact);
  return 0;

ERROR:
  return -1;
}

int decode_reply(const char *packet, size_t len, struct ClearviewInstance *cv,
                 struct ContactAddress *contact) {
  struct ReplyPacketHeader header;
  const char *end = packet + len;

  RETURN_IF_NEGATIVE(decode_reply_header(packet, end, &header));
  packet += REPLY_HEADER_SIZE;

  switch (header.action) {
  case ACT_REPLY_PING:
    return decode_ping_reply(&header, packet, end, cv, contact);
    break;

  case ACT_REPLY_FIND_NODE:
    return decode_find_node_reply(&header, packet, end, cv, contact);
    break;

  case ACT_REPLY_ERROR:
    return decode_error_reply(&header, packet, end, cv, contact);
    break;

  case ACT_REPLY_STORE: 
    // Ignore; we will never make a store request!
  default:
    fprintf(stderr, "Ignoring reply with action %d\n", header.action);
    break;
  }
  return 0;
}


static char* encode_contact_address(const struct ContactAddress *address,
                                    char *buf) {
  *((uint8_t *) buf) = 4;
  buf += 1;

  /* The address should already be in network byte order. */
  *((uint32_t *) buf) = address->addr.s_addr;
  buf += sizeof(address->addr.s_addr);

  *((uint16_t *) buf) = htons(address->port);
  buf += sizeof(address->port);

  return buf;
}


// WARNING: Not re-entrant!
// You tell it where to put the encoded header. You better get the size right.
void encode_request_header(const struct RequestPacketHeader *header,
                           char buf[REQUEST_HEADER_SIZE]) {
  char *start = buf;

  // WARNING: probably relies on your architecture doing unaligned stores.
  *((uint64_t *) buf) = bswap_64(header->connection_id);
  buf += sizeof(header->connection_id);

  *((uint32_t *) buf) = htonl(header->action);
  buf += sizeof(header->action);

  *((uint32_t *) buf) = htonl(header->transaction_id);
  buf += sizeof(header->transaction_id);

  *((uint8_t *) buf) = header->protocol_version;
  buf += sizeof(header->protocol_version);

  *((uint8_t *) buf) = header->vendor_id;
  buf += sizeof(header->vendor_id);

  *((uint32_t *) buf) = htonl(header->network_id);
  buf += sizeof(header->network_id);

  *((uint8_t *) buf) = header->local_protocol_version;
  buf += sizeof(header->local_protocol_version);

  buf = encode_contact_address(&header->node_address, buf);
  
  *((uint32_t *) buf) = htonl(header->instance_id);
  buf += sizeof(header->instance_id);

  *((uint64_t *) buf) = bswap_64(header->time);
  buf += sizeof(header->time);

  assert(buf == (start + REQUEST_HEADER_SIZE));
}


void encode_ping_request(const struct RequestPingPacket *packet, char **user_buf,
                         size_t *buflen) {
  static char buf[REQUEST_HEADER_SIZE];
  encode_request_header(&packet->header, buf);
  *user_buf = buf;
  *buflen = sizeof(buf);
}


void encode_find_node_request(const struct RequestFindNodePacket *packet,
                              char **user_buf, size_t *buflen) {
  static char buf[REQUEST_FIND_NODE_SIZE];
  char *p;

  encode_request_header(&packet->header, buf);
  p = buf + REQUEST_HEADER_SIZE;
  
  *((uint8_t *)p) = 20;
  p += sizeof(uint8_t);

  memcpy(p, packet->id, KEY_SIZE);
  p += KEY_SIZE;
  
  *((uint32_t *)p) = htonl(packet->node_status);
  p += sizeof(uint32_t);

  *((uint32_t *)p) = 0;
  p += sizeof(uint32_t);

  assert(p == (buf + REQUEST_FIND_NODE_SIZE));
  *user_buf = buf;
  *buflen = REQUEST_FIND_NODE_SIZE;
}


static void encode_reply_header(const struct ReplyPacketHeader *header,
                                char buf[REPLY_HEADER_SIZE]) {
  char *start = buf;

  (*(uint32_t *)buf) = htonl(header->action);
  buf += sizeof(header->action);

  (*(uint32_t *)buf) = htonl(header->transaction_id);
  buf += sizeof(header->transaction_id);

  (*(uint64_t *)buf) = bswap_64(header->connection_id);
  buf += sizeof(header->connection_id);

  (*(uint8_t *)buf) = header->protocol_version;
  buf += sizeof(header->protocol_version);

  (*(uint8_t *)buf) = header->vendor_id;
  buf += sizeof(header->vendor_id);

  (*(uint32_t *)buf) = htonl(header->network_id);
  buf += sizeof(header->network_id);

  (*(uint32_t *)buf) = htonl(header->instance_id);
  buf += sizeof(header->instance_id);

  assert(buf - start == REPLY_HEADER_SIZE);
}


static const char HARMLESS_VIVALDI[] = "\x01\x01\x10" // 1 position, type 1, sz16
  "\x00\x00\x00\x00" // x-coord
  "\x00\x00\x00\x00" // y-coord
  "\x00\x00\x00\x00" // height
  "\x7f\x80\x00\x00"; // error (+infty)


void encode_ping_reply(const struct ReplyPingPacket *packet, char **user_buf,
                       size_t *buflen) {
  static char buf[REPLY_HEADER_SIZE+sizeof(HARMLESS_VIVALDI)];
  encode_reply_header(&packet->header, buf);
  memcpy(buf+REPLY_HEADER_SIZE, HARMLESS_VIVALDI, sizeof(HARMLESS_VIVALDI));

  *user_buf = buf;
  *buflen = sizeof(buf);
}


// TODO(swolchok): support more than one contact. Fix callers to g_free the
// buf when you do that.
void encode_find_node_reply(const struct ReplyFindNodePacket *packet,
                            char **user_buf, size_t *buflen) {
  static char buf[REPLY_HEADER_SIZE + 4 + 4 + 4 + sizeof(HARMLESS_VIVALDI) +
                  2 + CONTACT_NET_SIZE];
  char *p = buf;

  g_assert(packet->contacts_count == 1);

  encode_reply_header(&packet->header, buf);
  p += REPLY_HEADER_SIZE;

  *((uint32_t *)p) = htonl(packet->spoof_id);
  p += sizeof(packet->spoof_id);

  *((uint32_t *)p) = htonl(packet->node_type);
  p += sizeof(packet->node_type);

  *((uint32_t *)p) = htonl(packet->dht_size);
  p += sizeof(packet->dht_size);

  memcpy(p, HARMLESS_VIVALDI, sizeof(HARMLESS_VIVALDI));
  p += sizeof(HARMLESS_VIVALDI);

  *((uint16_t*)p) = htons(packet->contacts_count);
  p += sizeof(packet->contacts_count);

  *((uint8_t*)p) = 1; // Contact type UDP.
  p += sizeof(uint8_t);

  *((uint8_t*)p) = SUPPORTED_PROTOCOL_VERSION;
  p += sizeof(uint8_t);

  p = encode_contact_address(packet->contacts[0], p);

  g_assert(p - buf == sizeof(buf));

  *user_buf = buf;
  *buflen = sizeof(buf);
}


void encode_store_reply(const struct ReplyStorePacket *packet,
                        char **user_buf, size_t *buflen) {
  static char buf[REPLY_HEADER_SIZE + 1];
  char *p = buf;

  // We don't actually support diversificiations.
  g_assert(packet->diversifications_length == 0);

  encode_reply_header(&packet->header, buf);
  p += REPLY_HEADER_SIZE;

  *p++ = '\0';

  g_assert(p - buf == sizeof(buf));
  
  *user_buf = buf;
  *buflen = sizeof(buf);
}
